package com.gitee.transformation

import org.apache.flink.api.common.functions.Partitioner
import org.apache.flink.streaming.api.scala.{DataStream, StreamExecutionEnvironment}

object PartitionTest {
  def main(args: Array[String]): Unit = {
    val senv: StreamExecutionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment
    senv.setParallelism(3)
    import org.apache.flink.api.scala._
    val dataStream: DataStream[(String, Int)] = senv.fromElements(("a", 3), ("d", 4), ("c", 2), ("c", 5), ("a", 5))
    // 操作数据
    //随机发送
    dataStream.shuffle.print("随机分区: ")
    //如果上游有两个分区,下游4个分区,上游一个分区给两个下游发送
    dataStream.rescale.print("本地轮询: ")
    //上游每个分区给下游每个分区发送
    dataStream.rebalance.print("轮询: ")

    dataStream.partitionCustom(new MyPartition(3),0).print("自定义的分区: ")
    senv.execute()
  }

  class MyPartition(numPartitions: Int) extends Partitioner[String] {
    override def partition(key: String, numPartitions: Int): Int = {
      key match {
        case "a" => 0
        case "b" => 1
        case _ => 2
      }
    }
  }

}
